# !/user/bin/env python3
# # # encoding: utf-8
'''
使用光流估计方法，在前述测试视频上计算特征点，进一步进行特征点光流估计。
'''
import cv2 as cv

videofileName = r'D:/vtest.avi'
#设置角点检测的参数
feature_params = dict(maxCorners=100, qualityLevel=0.3, minDistance=7, blockSize=7)
#设置lucas-Kanade光流法参数
lk_params = dict(winSize=(15, 15), maxLevel=2, criteria=(cv.TERM_CRITERIA_EPS|cv.TERM_CRITERIA_COUNT, 10, 0.03))
cap = cv.VideoCapture(videofileName)
#计算第一帧的特征点
ret, prev = cap.read()
prevGray = cv.cvtColor(prev, cv.COLOR_BGR2GRAY)
#筛选更符合要求的特征点
p0 = cv.goodFeaturesToTrack(prevGray, mask=None, **feature_params)

while True:
    ret, frame = cap.read()
    #如果没有读取到当前帧的图像，结束程序
    if not ret:
        break
    gray = cv.cvtColor(frame, cv.COLOR_BGR2GRAY)
    #计算光流
    p1, st, err = cv.calcOpticalFlowPyrLK(prevGray, gray, p0, None, **lk_params)
    #选取好的跟踪点
    goodPoints = p1[st == 1]
    goodPrePoints = p0[st == 1]
    #在结果图像中迭加画出特征点和计算出来的光流向量
    res = frame.copy()
    drawColor = (0, 0, 255)
    for i, (cur, prev) in enumerate(zip(goodPoints, goodPrePoints)):
        x0, y0 = cur.ravel()
        x1, y1 = prev.ravel()
        cv.line(res, (x0, y0), (x1, y1), drawColor)
        cv.circle(res, (x0, y0), 3, drawColor)
    #更新上一帧
    prevGray = gray.copy()
    p0 = goodPoints.reshape(-1, 1,2)
    #显示计算结果图像
    cv.imshow("Detect Result", res)
    #每帧间隔30ms
    key = cv.waitKey(30)
    #按下ESC终端程序并退出
    if key == 27:
        break


cap.release()
cv.destroyAllWindows()